# 2021.07.15 / author: kw-a
import numpy as np
import scipy.io as sio
import os

path = 'D:/data/dataSet/'  # .arff文件目录
save_path = 'D:/data/save1/'  # .mat文件保存目录

pathList = ['NASA']
# , 'Relink', 'SOFTLAB' 'AEEEM', 'MORPH',

# 主函数
def main():

    for dirName in pathList:
        files = os.listdir(path+dirName)

        for fileName in files:
            file_name = fileName.split('.')[0]
            if fileName == '.DS_Store':
                continue

            dataSet = []

            with open(path + dirName + '/' + file_name + '.arff') as fr:  # 打开读取文件
                Iter = iter(fr.readlines())  # 创建迭代器
                curline = ['0']  # 初始化
                # 迭代到@data结束
                while curline[0] != '@data':
                    line = next(Iter)  # 迭代
                    curline = line.strip().split(',')
                    print(curline)
                print('开始转换{}数据'.format(fileName))
                while curline[0]:
                    try:
                        line = next(Iter)
                        # 跳过空行
                        if line == '\n':
                            line = next(Iter)
                        curline = line.strip().split(',')  # 去掉首尾空格或换行符
                        attrNum = len(curline) - 1  # 属性个数
                        fltLine = list(map(float, curline[0:attrNum]))  # 属性值，从下标0读到下标2，不包括下标3，即3个元素，将其转换为浮点数

                        # if curline[attrNum] == 'buggy':
                        if curline[attrNum] == 'Y':
                            label = int(1)
                        else:
                            label = int(-1)

                        fltLine.append(label)
                        dataSet.append(fltLine)
                    except StopIteration:  # next(Iter)迭代结束
                        break
            data_array = np.array(dataSet)
            sio.savemat(save_path + dirName + '/' + file_name + '.mat', mdict={file_name: data_array})  # 转换为.mat文件

    # arff_to_mat_one('EQ', 'EQ')  # filename, mat_name
    # arff_to_mat_two('Forest Covertype', 'FCT')

    # 查看.mat数据
    # data = sio.loadmat(save_path + 'EQ' + '.mat')
    # print('FCT: ', data)

'''
方法一
读取所有行数据readlines()
创建迭代器进行迭代
'''
def arff_to_mat_one(filename, mat_name):
    dataSet = []
    with open(path + filename + '.arff') as fr:  # 打开读取文件
        Iter = iter(fr.readlines())  # 创建迭代器
        curline = ['0']  # 初始化
        # 迭代到@data结束
        while curline[0] != '@data':
            line = next(Iter)  # 迭代
            curline = line.strip().split(',')
            print(curline)
        print('开始转换{}数据'.format(filename))
        while curline[0]:
            try:
                line = next(Iter)
                # 跳过空行
                if line == '\n':
                    line = next(Iter)
                curline = line.strip().split(',')  # 去掉首尾空格或换行符
                attrNum = len(curline) - 1  # 属性个数
                fltLine = list(map(float, curline[0:attrNum]))  # 属性值，从下标0读到下标2，不包括下标3，即3个元素，将其转换为浮点数

                # print(curline[attrNum])

                if curline[attrNum] == 'buggy':
                    label = int(1)
                else:
                    label = int(-1)

                # label = int((curline[attrNum]))  # 标签 添加最后一个元素
                fltLine.append(label)
                dataSet.append(fltLine)
            except StopIteration:  # next(Iter)迭代结束
                break
    data_array = np.array(dataSet)
    sio.savemat(save_path + mat_name + '.mat', mdict={mat_name: data_array})  # 转换为.mat文件


# '''
# 方法二
# 逐行读取readline()
# '''
# def arff_to_mat_two(filename, mat_name):
#     dataSet = []
#     with open(path + filename + '.arff') as fr:  # 打开读取文件
#         curline = ['0']  # 初始化
#         # 迭代到@data结束
#         while curline[0] != '@data':
#             line = fr.readline()  # 读取一行数据
#             curline = line.strip().split(',')
#             print(curline)
#         print('开始转换{}数据'.format(filename))
#         while curline[0]:
#             try:
#                 line = fr.readline()
#                 if line == '\n':  # 跳过空行
# 					line = fr.readline()
#                 curline = line.strip().split(',')  # 去掉首尾空格或换行符
#                 attrNum = len(curline) - 1  # 属性个数
#                 fltLine = list(map(float, curline[0:attrNum]))  # 属性值，从下标0读到下标2，不包括下标3，即3个元素，将其转换为浮点数
#                 label = int((curline[attrNum]))  # 标签 添加最后一个元素
#                 fltLine.append(label)
#                 dataSet.append(fltLine)
#             except ValueError:  # 迭代结束
#                 break
#     data_array = np.array(dataSet)
#     sio.savemat(save_path + mat_name + '.mat', mdict={mat_name: data_array})  # 转换为.mat文件


if __name__ == '__main__':
    main()
